import time
import copy
import config as cfg
from Scripts.sensor import Sensor
import warnings



class PID:
    def __init__(
        self, 
        kp: float, ki: float, kd: float, 
        sample_time: float = 0.02, 
        output_range_low: float = None, output_range_high: float = None
    ):
        
        self.kp = kp
        self.ki = ki
        self.kd = kd
        
        self.sample_time = sample_time

        self.previous_error = 0.0
        self.integral = 0.0
        self.last_time = None
        self.current_time = None
        
        self.output_range_low = output_range_low
        self.output_range_high = output_range_high

    def reset(self):
        self.previous_error = 0.0
        self.integral = 0.0
        self.last_time = None
        self.current_time = None

    def update(self, current_value: float, setpoint: float):
        """Compute and return the PID control output."""
        error = setpoint - current_value
        # Handle time delta
        self.current_time = time.time()
        delta_time = self.current_time - self.last_time if self.last_time is not None else self.sample_time
        self.last_time = copy.deepcopy(self.current_time)
        # Proportional term
        proportional = self.kp * error
        # Integral term
        self.integral += error * delta_time
        integral = self.ki * self.integral
        # Derivative term
        derivative = self.kd * (error - self.previous_error) / delta_time if delta_time > 0 else 0.0
        # Save state for next update
        self.previous_error = error
        # Control output
        output = proportional + integral + derivative
        if self.output_range_low == None or self.output_range_high == None:
            return output
        else:
            return max(min(output, self.output_range_high), self.output_range_low)

class WheelVelocityController:
    def __init__(self):
        self.wheel_vel_controller = []
        for wheel_name in cfg.WHEEL["names"]:
            self.wheel_vel_controller.append(
                PID(
                    kp = cfg.WHEEL["wheel_vel_controller"]["kp"],
                    ki = 0,
                    kd = cfg.WHEEL["wheel_vel_controller"]["kd"],
                    output_range_low = -abs(cfg.WHEEL["wheel_vel_controller"]["max_torque_change"]),
                    output_range_high = abs(cfg.WHEEL["wheel_vel_controller"]["max_torque_change"])
                )
            )
    
    # 根据当前帧获取actuator的驱动力数据
    def get_torque(self, pos_sensor_data: list, target_pos: list) -> list:
        actuator_torque = []
        for index in range(0, len(target_pos)):
            output = self.wheel_vel_controller[index].update(current_value = pos_sensor_data[index], setpoint = target_pos[index])
            actuator_torque.append(output)
        return actuator_torque
    

class Controller:
    def __init__(self, model, data):
        self.model = model
        self.data = data
        self.sensor = Sensor(self.model, self.data)
        self.wheelVelocityController = WheelVelocityController()

        self.wheel_torque = [0.0] * 4
        self.wheel_vel_target = [0.0] * 4
        self.wheel_LR_error = 0
        
        self.wheel_LR_error_controller = PID(
            kp = cfg.WHEEL["wheel_FL_error_controller"]["kp"],
            ki = 0,
            kd = cfg.WHEEL["wheel_FL_error_controller"]["kd"],
            output_range_low = -abs(cfg.WHEEL["wheel_FL_error_controller"]["max_torque_change"]),
            output_range_high = abs(cfg.WHEEL["wheel_FL_error_controller"]["max_torque_change"])
        )
        
    # 获取轮子的力矩
    def __get_wheel_torque(self, wheel_vel_target: list ) -> list:
        if len(wheel_vel_target) != 4:
            warnings.warn(f"输入的wheel_vel_target长度不正确, 输入的wheel_vel_target: {wheel_vel_target}", UserWarning, stacklevel=1)
        else:
            wheel_vel_now = self.sensor.get_wheel_vel_sensor_data()
            wheel_torque_change = self.wheelVelocityController.get_torque(wheel_vel_now, wheel_vel_target)
            for i in range(len(wheel_torque_change)):
                self.wheel_torque[i] += wheel_torque_change[i]
                self.wheel_torque[i] = max(
                    min(self.wheel_torque[i], cfg.WHEEL["wheel_vel_controller"]["torque_range"][1]), 
                    cfg.WHEEL["wheel_vel_controller"]["torque_range"][0]
                )
        return self.wheel_torque

    # 轮速控制
    def car_vel_controller(self, target_vel):
        R = 0.05 # 轮子半径
        C = 0.34 # 轴距

        V_l = target_vel[0] - target_vel[1] * C / 2
        V_r = target_vel[0] + target_vel[1] * C / 2

        self.wheel_vel_target = [V_l / R, V_r / R, V_l / R, V_r / R]
        
        wheel_torque = self.__get_wheel_torque(self.wheel_vel_target)
        
        wheel_torque = self.__rotation_compensation(target_vel, wheel_torque)

        return wheel_torque
    
    # 转动补偿
    def __rotation_compensation(self, target_vel, wheel_torque):
        now_angular_vel = self.sensor.get_imu_sensor_data()[1][2]
        wheel_LR_error_change = self.wheel_LR_error_controller.update(now_angular_vel, target_vel[1])
        self.wheel_LR_error += wheel_LR_error_change
        self.wheel_LR_error = max(
                min(self.wheel_LR_error, cfg.WHEEL["wheel_FL_error_controller"]["torque_range"][1]), 
                cfg.WHEEL["wheel_FL_error_controller"]["torque_range"][0]
        )
        for i in range(4):
            if i % 2 == 0:
                wheel_torque[i] -= self.wheel_LR_error
            else:
                wheel_torque[i] += self.wheel_LR_error
        return wheel_torque
    
    # 获取轮速目标
    def get_wheel_vel_target(self):
        return self.wheel_vel_target

        
        